package com.zrc.japktool.util;

import net.lingala.zip4j.ZipFile;
import net.lingala.zip4j.model.FileHeader;
import net.lingala.zip4j.model.ZipParameters;
import net.lingala.zip4j.model.enums.CompressionLevel;
import net.lingala.zip4j.model.enums.CompressionMethod;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class ApkZipalign {

    private static final int ALIGNMENT = 4;
    private static final int THREAD_POOL_SIZE = Math.max(4, Runtime.getRuntime().availableProcessors());
    private static final int MAX_QUEUE_SIZE = 1000;
    private static final int LOG_PROGRESS_INTERVAL = 100; // 每处理100个文件输出一次进度

    public static void alignApk(String inputApkPath, String outputApkPath, TaskLog task) throws Exception {
        long startTime = System.currentTimeMillis();
        task.tip("开始APK对齐处理: " + inputApkPath);

        if (!Files.exists(Paths.get(inputApkPath))) {
            throw new FileNotFoundException("输入APK文件不存在: " + inputApkPath);
        }

        // 统计信息
        AtomicInteger totalFiles = new AtomicInteger(0);
        AtomicInteger alignedFiles = new AtomicInteger(0);
        AtomicInteger processedFiles = new AtomicInteger(0);

        // 使用更高效的线程池配置
        ExecutorService executor = new ThreadPoolExecutor(
                THREAD_POOL_SIZE,
                THREAD_POOL_SIZE,
                30L, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(MAX_QUEUE_SIZE),
                new ThreadPoolExecutor.CallerRunsPolicy());

        Path tempDir = Files.createTempDirectory("apk_align");
        try {
            // 1. 解压APK
            task.info("解压APK文件...");
            long extractStart = System.currentTimeMillis();
            Path extractDir = tempDir.resolve("extract");
            Files.createDirectory(extractDir);

            try (ZipFile zipFile = new ZipFile(inputApkPath)) {
                zipFile.setRunInThread(false);
                zipFile.extractAll(extractDir.toString());
            }
            task.info(String.format("解压完成, 耗时: %dms", System.currentTimeMillis() - extractStart));

            // 2. 收集文件信息
            task.info("扫描文件...");
            List<Path> allFiles = new ArrayList<>();
            List<Path> filesToAlign = new ArrayList<>();

            try (var paths = Files.walk(extractDir)) {
                paths.filter(Files::isRegularFile).forEach(path -> {
                    allFiles.add(path);
                    try {
                        String fileName = extractDir.relativize(path).toString().replace("\\", "/");
                        long size = Files.size(path);
                        if (!fileName.equals("resources.arsc") &&
                                !fileName.equals("AndroidManifest.xml") &&
                                !fileName.equals("classes0.jar") &&
                                !fileName.equals("libSecShell32") &&
                                !fileName.endsWith(".so") &&
                                size > 0 &&
                                (size & (ALIGNMENT - 1)) != 0) {
                            filesToAlign.add(path);
                        }
                    } catch (IOException e) {
                        task.error("处理文件 " + path + " 失败: " + e.getMessage());
                    }
                });
            }

            totalFiles.set(allFiles.size());
            task.info(String.format("共发现 %d 个文件, 其中 %d 个需要对齐", totalFiles.get(), filesToAlign.size()));

            // 3. 并行处理文件
            task.info("开始处理文件...");
            long processStart = System.currentTimeMillis();
            List<Future<?>> futures = new ArrayList<>();

            // 优先处理关键文件
            Path resourcesPath = extractDir.resolve("resources.arsc");
            if (!Files.exists(resourcesPath)) {
                throw new IOException("APK中未找到resources.arsc");
            }
            futures.add(executor.submit(() -> {
                processResourcesArsc(resourcesPath, task);
                processedFiles.incrementAndGet();
            }));

            Path manifestPath = extractDir.resolve("AndroidManifest.xml");
            if (Files.exists(manifestPath)) {
                futures.add(executor.submit(() -> {
                    alignFile(manifestPath, task);
                    processedFiles.incrementAndGet();
                    alignedFiles.incrementAndGet();
                }));
            }

            // 处理其他需要对齐的文件
            for (Path file : filesToAlign) {
                futures.add(executor.submit(() -> {
                    alignFile(file, task);
                    processedFiles.incrementAndGet();
                    alignedFiles.incrementAndGet();

                    // 定期输出进度
                    if (processedFiles.get() % LOG_PROGRESS_INTERVAL == 0) {
                        task.info(String.format("处理进度: %d/%d (已对齐: %d)", processedFiles.get(), totalFiles.get(), alignedFiles.get()));
                    }
                }));
            }

            // 等待所有任务完成
            for (Future<?> future : futures) {
                try {
                    future.get();
                } catch (ExecutionException e) {
                    throw new IOException("文件处理失败", e.getCause());
                }
            }
            task.info(String.format("文件处理完成, 耗时: %dms", System.currentTimeMillis() - processStart));
            task.info(String.format("总计: 处理 %d 个文件, 对齐 %d 个文件", processedFiles.get(), alignedFiles.get()));

            // 4. 重新打包APK
            task.info("开始重新打包APK...");
            long zipStart = System.currentTimeMillis();
            try (ZipFile newZip = new ZipFile(outputApkPath)) {
                newZip.setRunInThread(false);

                // 特殊处理resources.arsc
                ZipParameters resourcesParams = new ZipParameters();
                resourcesParams.setFileNameInZip("resources.arsc");
                resourcesParams.setCompressionMethod(CompressionMethod.STORE);
                resourcesParams.setCompressionLevel(CompressionLevel.NORMAL);
                newZip.addFile(resourcesPath.toFile(), resourcesParams);

                // 添加其他文件
                AtomicInteger zippedFiles = new AtomicInteger(0);
                for (Path file : allFiles) {
                    if (!file.equals(resourcesPath)) {
                        String entryName = extractDir.relativize(file).toString().replace("\\", "/");
                        ZipParameters params = new ZipParameters();
                        params.setFileNameInZip(entryName);
                        params.setCompressionMethod(entryName.equals("AndroidManifest.xml") ? CompressionMethod.STORE : CompressionMethod.DEFLATE);
                        params.setCompressionLevel(CompressionLevel.NORMAL);
                        newZip.addFile(file.toFile(), params);

                        zippedFiles.incrementAndGet();
                        if (zippedFiles.get() % LOG_PROGRESS_INTERVAL == 0) {
                            task.info(String.format("打包进度: %d/%d", zippedFiles.get(), totalFiles.get()));
                        }
                    }
                }
            }
            task.info(String.format("打包完成, 耗时: %dms", System.currentTimeMillis() - zipStart));

            // 5. 验证
            task.info("验证结果...");
            verifyResourcesArsc(Paths.get(outputApkPath).toFile(), task);

            long totalTime = System.currentTimeMillis() - startTime;
            task.info(String.format("APK对齐完成! 总耗时: %dms (%.2fs)", totalTime, totalTime/1000.0));
            task.info("输出文件: " + outputApkPath);

        } finally {
            // 清理资源
            task.info("清理临时文件...");
            executor.shutdown();
            try {
                if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                    executor.shutdownNow();
                }
            } catch (InterruptedException e) {
                executor.shutdownNow();
                Thread.currentThread().interrupt();
            }
            deleteDir(tempDir.toFile());
        }
    }

    private static void processResourcesArsc(Path resourcePath, TaskLog task) {
        try {
            task.info("处理resources.arsc (强制未压缩+4字节对齐)");
            byte[] originalData = Files.readAllBytes(resourcePath);
            int alignedLength = (originalData.length + ALIGNMENT - 1) & ~(ALIGNMENT - 1);
            if (alignedLength != originalData.length) {
                byte[] alignedData = Arrays.copyOf(originalData, alignedLength);
                Files.write(resourcePath, alignedData);
                task.info(String.format("resources.arsc 对齐: %d -> %d 字节", originalData.length, alignedLength));
            }
        } catch (Exception e) {
            task.error("处理resources.arsc失败: " + e.getMessage());
            throw new RuntimeException(e);
        }
    }

    private static void alignFile(Path filePath, TaskLog task) {
        try {
            String fileName = filePath.getFileName().toString();
            byte[] originalData = Files.readAllBytes(filePath);
            int originalLength = originalData.length;
            int alignedLength = (originalLength + ALIGNMENT - 1) & ~(ALIGNMENT - 1);

            if (alignedLength == originalLength) {
                return;
            }

            byte[] alignedData = Arrays.copyOf(originalData, alignedLength);
            Files.write(filePath, alignedData);
            task.info(String.format("对齐文件: %s (%d -> %d 字节)", fileName, originalLength, alignedLength));
        } catch (Exception e) {
            task.error("对齐文件失败: " + filePath + " - " + e.getMessage());
            throw new RuntimeException(e);
        }
    }

    private static void verifyResourcesArsc(File apkFile, TaskLog task) throws IOException {
        try (ZipFile zipFile = new ZipFile(apkFile)) {
            FileHeader header = zipFile.getFileHeader("resources.arsc");
            if (header == null) {
                throw new IOException("验证失败：resources.arsc丢失");
            }

            if (header.getCompressionMethod() != CompressionMethod.STORE) {
                throw new IOException("验证失败：resources.arsc仍被压缩");
            }

            long uncompressedSize = header.getUncompressedSize();
            if ((uncompressedSize & (ALIGNMENT - 1)) != 0) {
                throw new IOException("验证失败：resources.arsc未4字节对齐");
            }

            task.info("验证通过: resources.arsc (未压缩, " + uncompressedSize + " 字节, 已对齐)");
        }
    }

    private static boolean deleteDir(File dir) {
        if (dir.isDirectory()) {
            File[] children = dir.listFiles();
            if (children != null) {
                for (File child : children) {
                    deleteDir(child);
                }
            }
        }
        return dir.delete();
    }

    public static void main(String[] args) throws Exception {
//        String inputPath = "C:/Users/15666/Desktop/test/merged_1756797042034.apk";
//        String outputPath = "C:/Users/15666/Desktop/test/demo_aligned.apk";
//
//        try {
//            alignApk(inputPath, outputPath, System.out::println);
//
//            // 签名步骤
//            String sign_inApkPath = outputPath;
//            String outApkDirPath = new File(outputPath).getParent().replace("\\", "/");
//            String sign_outApkName = new File(outputPath).getName().replace(".apk", "");
//            String sign_outApkPath = new File(outApkDirPath, sign_outApkName + "-zipalign-sign.apk").getAbsolutePath().replace("\\", "/");
//
//            System.out.println("开始签名APK...");
//            long signStart = System.currentTimeMillis();
//            ApkSign.sign(
//                    "android",
//                    "androiddebugkey",
//                    "android",
//                    false,
//                    Paths.get(sign_inApkPath),
//                    Paths.get(sign_outApkPath),
//                    System.out::println);
//            System.out.printf("APK签名完成! 耗时: %dms%n", System.currentTimeMillis() - signStart);
//            System.out.println("签名后文件: " + sign_outApkPath);
//
//            if (Paths.get(outputPath).toFile().exists()) {
//                Files.delete(Paths.get(outputPath));
//            }
//        } catch (Exception e) {
//            System.err.println("处理APK时出错: " + e.getMessage());
//            e.printStackTrace();
//            System.exit(1);
//        }
    }
}